import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

plt.rcParams['font.sans-serif'] = ['STHeiti', 'PingFang SC', 'Heiti TC', 'Microsoft YaHei', 'SimHei', 'Arial Unicode MS']  
plt.rcParams['axes.unicode_minus'] = False

fig, axs = plt.subplots(1, 3, figsize=(12, 4))  # 改为一行三列图

species = (
    "2025",
    "2030",
    "2035",
    "2040",
    "2045",
    "2050",
)

def read_csv(csv_path):
    df = pd.read_csv(csv_path, index_col=0)
    df.rename(index={"Hydro_Nonpumped": "Hydro", "Central_PV": "PV"}, inplace=True)
    df.columns = species

    energy = ["Coal", "Gas", "Nuclear", "Hydro", "PV", "Wind"]
    weights = {row: np.array(df.loc[row].values) for row in energy}

    # 按列归一化：每列除以该列的总和
    normalized_df = df.div(df.sum(axis=0), axis=1)

    # 转成字典形式，每行是一个 np.array
    normalized_weights = {
        row: normalized_df.loc[row].values
        for row in energy
    }

    return df, weights, normalized_weights

width = 0.8
colors = ['#000000', '#ffdbba', '#f4a261', '#00c0ff', '#ffd800', '#39b371']

def plot_in_axes(axes, species, weights):
    bottom = np.zeros(6)
    for i, (boolean, weight) in enumerate(weights.items()):
        axes.bar(species, weight, width, label=boolean, bottom=bottom, color=colors[i])
        bottom += weight

def plot_csv_in_col(csv_path, col):
    df, w, nw = read_csv(csv_path)

    ax = axs[col]
    ax.set_ylim(top=15000)
    ax.set_yticks(np.arange(0, 15001, 2500))  # 修改为每 2500 一格
    ax.grid(axis='y', linestyle='-', alpha=0.6)

    plot_in_axes(ax, species, w)

    # 添加第二个 Y 轴
    ax2 = ax.twinx()
    share = nw["PV"] + nw["Wind"]
    line = ax2.plot(species, share, color='blue', marker='o', label='PV+Wind Share')
    ax2.set_ylim(0, 0.6)
    ax2.set_yticks(np.arange(0, 0.61, 0.1))  # 修改为每 0.1 一格
    ax2.set_ylabel("PV+Wind Share")
    ax2.tick_params(axis='y')
    return line[0]  # 返回折线对象

# 绘制三个子图并获取折线对象
line1 = plot_csv_in_col("data-10fix.csv", 0)
line2 = plot_csv_in_col("data-10v1g.csv", 1)
line3 = plot_csv_in_col("data-10v2g.csv", 2)

# 设置图例在整个图的顶部
handles, labels = axs[0].get_legend_handles_labels()
fig.legend(handles + [line1], labels + ['PV+Wind Share'], loc='upper center', ncol=7, bbox_to_anchor=(0.5, 1))

axs[0].set_ylabel("Power Generation (TWh)")
axs[0].set_xlabel("随到随充")
axs[1].set_ylabel("Power Generation (TWh)")
axs[1].set_xlabel("有序充电")
axs[2].set_xlabel("车网互动")
axs[2].set_ylabel("Power Generation (TWh)")

plt.tight_layout(rect=[0, 0, 1, 0.9])
plt.show()
